import torch
import torch.nn as nn


class Model(nn.Module):
    def __init__(self, actor_input_size, actor_output_size, actor_hidden_size,
                 critic_input_size, critic_output_size, critic_hidden_size):
        super(Model, self).__init__()
        self.actor_model = ActorModel(actor_input_size, actor_output_size, actor_hidden_size)
        self.critic_model = CriticModel(critic_input_size, critic_output_size, critic_hidden_size)

    def policy(self, obs):
        return self.actor_model.policy(obs)

    def value(self, obs, act):
        return self.critic_model.value(obs, act)


class ActorModel(nn.Module):
    def __init__(self, input_size, output_size, hidden_size=100):
        super(ActorModel, self).__init__()
        self.dense1 = nn.Linear(input_size, hidden_size)
        self.dense2 = nn.Linear(hidden_size, output_size)
        self.relu = nn.ReLU()
        self.tanh = nn.Tanh()

    def policy(self, obs):
        hid = self.relu(self.dense1(obs))
        out = self.tanh(self.dense2(hid))
        return out


class CriticModel(nn.Module):
    def __init__(self, input_size, output_size=1, hidden_size=100):
        super(CriticModel, self).__init__()
        self.dense1 = nn.Linear(input_size, hidden_size)
        self.dense2 = nn.Linear(hidden_size, output_size)
        self.relu = nn.ReLU()

    def value(self, obs, act):
        concat = torch.cat([obs, act], dim=1)
        hid = self.relu(self.dense1(concat))
        Q = self.dense2(hid)
        Q = Q.squeeze(1)
        return Q


if __name__ == '__main__':
    model = Model(actor_input_size=4, actor_output_size=6, actor_hidden_size=100,
                  critic_input_size=5, critic_output_size=1, critic_hidden_size=100)

    a = model.policy(torch.randn((3, 4)))
    b = model.value(torch.randn((3, 4)), torch.randn((3, 1)))
    c = model.parameters()
    d = model.state_dict()
    print(c)
    tau = 0.05
    model_state = model.state_dict().copy()
    target_model_state = model.state_dict().copy()
    for s in model_state:
        target_model_state[s] = tau * model_state[s] + (0.1 - tau) * target_model_state[s]
    print(target_model_state)
    model.load_state_dict(target_model_state)
    print(model)
